from argparse import ArgumentParser
from distutils.ccompiler import gen_lib_options
from random import randint
import csv

''' 
生成地图到文件: maze_map.csv
参数: --size, -s : 地图大小，矩形地图
grid = [[0 for x in range(width)] for y in range(height)]
'''


def genMap(size, filename) -> None:

    maze_map = [[0 for x in range(size)] for y in range(size)]
    #  随机生成地图
    for x in range(size):
        for y in range(size):
            maze_map[x][y] = randint(2, 4) % 2
    #  设置起点，终点
    maze_map[size-1][size-1] = 3
    maze_map[0][0] = 2

    # 把maze_map写入文件
    with open(filename, 'w') as f:
        f_csv = csv.writer(f)
        f_csv.writerows(maze_map)


def genQtable(size, filename) -> None:
    qtable = [[0 for x in range(4)] for y in range(size)]
    with open(filename, 'w') as f:
        f_csv = csv.writer(f)
        f_csv.writerows(qtable)


def genRewardTable(maze_map, filename, road_value, wall_value, dest_value) -> None:
    size = len(maze_map)
    reward = [[0 for x in range(size)] for y in range(size)]
    for i in range(size):
        for j in range(size):
            if maze_map[i][j] == '0' or maze_map[i][j] == '2':
                reward[i][j] = road_value
            elif maze_map[i][j] == '1':
                reward[i][j] = wall_value
            elif maze_map[i][j] == '3':
                reward[i][j] = dest_value
    with open(filename, 'w') as f:
        f_csv = csv.writer(f)
        f_csv.writerows(reward)


if __name__ == "__main__":

    # 从命令行得到size参数
    argparser = ArgumentParser(description="genmap")
    argparser.add_argument(
        '--size', '-s', help='set map size: n * n, max value 20.')
    result = argparser.parse_args()

    # size:地图大小
    size = int(result.size)

    if size > 20:
        size = 20

    genMap(size, 'maze_map.csv')
    genQtable(size * size, 'qtable.csv')

    maze_map = []
    with open('maze_map.csv', 'r') as f:
        f_csv = csv.reader(f)
        for row in f_csv:
            maze_map.append(row)

    genRewardTable(maze_map, 'reward1.csv', -0.1, -0.85, 10)
    genRewardTable(maze_map, 'reward2.csv', -0.3, -0.9, 100)
